package cn.genmer.test.security.machinelearning.deeplearning4j.mnist.V1;

import org.deeplearning4j.datasets.iterator.impl.MnistDataSetIterator;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.inputs.InputType;
import org.deeplearning4j.nn.conf.layers.ConvolutionLayer;
import org.deeplearning4j.nn.conf.layers.DenseLayer;
import org.deeplearning4j.nn.conf.layers.OutputLayer;
import org.deeplearning4j.nn.conf.layers.SubsamplingLayer;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.optimize.listeners.ScoreIterationListener;
import org.deeplearning4j.util.ModelSerializer;
import org.nd4j.evaluation.classification.Evaluation;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.learning.config.Adam;
import org.nd4j.linalg.lossfunctions.LossFunctions;

import java.io.File;

/**
 * 模型训练 数据集使用DL4j内置 也有网页版本: http://yann.lecun.com/exdb/mnist/
 */
public class MnistTrain {
    public static final String BASE_PATH = "/Users/genmer/Documents/Codes/tensorFlowModel/mnist";

    private int i;
    public static void train() throws Exception {

        int batchSize = 32;
        int numEpochs = 5;
        int seed = 12345;
        // 设置训练数据和测试数据迭代器
        DataSetIterator trainIter = new MnistDataSetIterator(batchSize, true, seed);
        DataSetIterator testIter = new MnistDataSetIterator(batchSize, false, seed);


        // 构建神经网络的配置
        MultiLayerNetwork model = originModel();
        System.out.println("模型配置信息：" + model.getLayerWiseConfigurations());
        model.init();
 
        model.setListeners(new ScoreIterationListener(1));
        // 训练模型
        for (int i = 0; i < numEpochs; i++) {
            model.fit(trainIter);
        }
 
        // 在测试数据上评估模型
        Evaluation eval = model.evaluate(testIter);
        System.out.println(eval.stats());
 
        // 将模型保存到本地磁盘
        File locationToSave = new File(BASE_PATH+"/cnn_mnist_model.zip");
        boolean saveUpdater = true; // 是否保存updater（用于进行模型参数更新）
        ModelSerializer.writeModel(model, locationToSave, saveUpdater);
    }

    /**
     * 模型构建（之前大坑是使用了L2正则）
     * @return
     */
    public static MultiLayerNetwork originModel(){
        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder()
                .seed(12345)
                .updater(new Adam(0.001))
                .list()
                .layer(0, new ConvolutionLayer.Builder(5, 5)
                        .nIn(1)
                        .stride(1, 1)
                        .nOut(20)
                        .activation(Activation.IDENTITY)
                        .build())
                .layer(1, new SubsamplingLayer.Builder(SubsamplingLayer.PoolingType.MAX)
                        .kernelSize(2, 2)
                        .stride(2, 2)
                        .build())
                .layer(2, new ConvolutionLayer.Builder(5, 5)
                        .stride(1, 1)
                        .nOut(50)
                        .activation(Activation.IDENTITY)
                        .build())
                .layer(3, new SubsamplingLayer.Builder(SubsamplingLayer.PoolingType.MAX)
                        .kernelSize(2, 2)
                        .stride(2, 2)
                        .build())
                .layer(4, new DenseLayer.Builder().nOut(500).build())
                .layer(5, new OutputLayer.Builder(LossFunctions.LossFunction.NEGATIVELOGLIKELIHOOD)
                        .nOut(10)
                        .activation(Activation.SOFTMAX)
                        .build())
                .setInputType(InputType.convolutionalFlat(28, 28, 1))
                .build();

        // 初始化模型并设置参数
        return new MultiLayerNetwork(conf);
    }

    public static void main(String[] args) throws Exception {
        train();
    }
}